fix(pd): handle border scalars and CPU tensors - #5832
Conversation
Store nlocal and nghost in one-element host tensors even when there are no swaps, and select local forward/backward copy primitives from the actual Paddle tensor place. Add direct custom-op regressions because existing Paddle model tests did not exercise nswap == 0 or CPU self-swaps in CUDA-enabled builds. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe Paddle border operation now stages scalar controls in single-element tensors, selects local-copy operations from tensor placement, and updates backward gradient handling. Tests cover zero swaps and CPU self-copy behavior. ChangesPaddle border operation fixes
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pd/test_border_op.py (1)
44-81: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winIncorrect gradient expectation in test.
The backward pass of
border_oproutes the gradient from the ghost slots back to their owner's local slot. In this test, index 2 (the ghost slot) receives its value from index 1 (the owner).
Duringresult.sum().backward(), the gradient for all elements ofresultis1.0. The backward operation accumulates the ghost gradient into the local owner usingindex_add_. Therefore, the gradient at index 1 should be1.0 + 1.0 = 2.0, while indices 0 and 2 retain their original gradients of1.0.The assertion expects
np.ones([3, 2]), which would incorrectly mean all gradients are1.0, causing the test to fail.💚 Proposed fix
- np.testing.assert_array_equal(g1_leaf.grad.numpy(), np.ones([3, 2])) + expected_grad = np.ones([3, 2]) + expected_grad[1] = 2.0 + np.testing.assert_array_equal(g1_leaf.grad.numpy(), expected_grad)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pd/test_border_op.py` around lines 44 - 81, Update the gradient assertion in test_border_op_self_copy_uses_cpu_place to expect accumulated gradients: index 1 should contain 2.0 in both columns, while indices 0 and 2 remain 1.0. Keep the forward result assertion and backward invocation unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/op/pd/comm.cc`:
- Around line 15-37: Wrap the definition of copy_local_tensor_data in `#if`
defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) and the matching `#endif`,
aligning it with its guarded callers so CPU-only builds do not compile
references to gpuMemcpy or gpuMemcpyDeviceToDevice.
---
Outside diff comments:
In `@source/tests/pd/test_border_op.py`:
- Around line 44-81: Update the gradient assertion in
test_border_op_self_copy_uses_cpu_place to expect accumulated gradients: index 1
should contain 2.0 in both columns, while indices 0 and 2 remain 1.0. Keep the
forward result assertion and backward invocation unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 17d433ad-098c-4824-b749-a97e187fecbb
📒 Files selected for processing (2)
source/op/pd/comm.ccsource/tests/pd/test_border_op.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5832 +/- ##
==========================================
+ Coverage 78.57% 79.35% +0.78%
==========================================
Files 1049 1085 +36
Lines 120659 126409 +5750
Branches 4349 4599 +250
==========================================
+ Hits 94807 100316 +5509
- Misses 24288 24438 +150
- Partials 1564 1655 +91 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both fixes are correct and this is a well-scoped port — pt/comm.cc does not share either bug (it already reads the scalars with .item() and dispatches the self-copy on is_cuda()), so this is pd catching up rather than an incomplete cross-backend fix. The {1} scalar shape and the phi::is_gpu_place-based copy_local_tensor_data are both right. One coverage note is inline; one out-of-scope parity gap for a follow-up:
pd backward lacks pt's ghost-row gradient zeroing (pre-existing, not this PR's scope). In Border_backward_t, d_local_g1_tensor is initialized as a copy of the incoming gradient (
deepmd-kit/source/op/pd/comm.cc
Lines 262 to 266 in 00eccf4
index_adds received gradient into owner rows — nothing zeroes the ghost range. pt/comm.cc does this explicitly (d_local_g1_tensor.slice(0, nlocal, ntotal).zero_()) because otherwise the ghost rows retain the raw incoming dL/dg_out[ghost], which is spurious whenever an exchanged feature is consumed as a genuine per-node leaf downstream (attention node embeddings in dpa2/dpa3, spin). pd and pt have diverged here — worth a separate follow-up (and while there, double-check the plain-CPU-build writeback path of the computed d_local_g1).
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The paddle.set_device("cpu") call is the right change and it fixes the half of my point that was fixable in this file -- on a paddlepaddle-gpu build the control and data tensors now land on CPU, so is_gpu_place() returns false and the branch matches the test's name.
The other half turns out to be worse than I described, and I got a detail wrong myself, so let me set it out properly. Inline below.
Approving anyway: the scalar handling and the CPU copy path read correctly, the change is small and self-contained, and the coverage problem is an infrastructure fact rather than something wrong with this diff. But I would not treat the green checks as evidence that this code works.
No CI job currently builds Paddle with CUDA, so the CPU-branch regression cannot fail in any pipeline. Document that intent in the test. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed head d4843dd96cf4f388e348140cb7704af35bf2b689 with three independent full subagent reviews. Two P1 backward-pass defects remain and are documented inline: reverse communication is discarded on most build paths, and the new test asserts the resulting incorrect gradient while ghost-input gradients remain nonzero. Because this PR was opened by the active njzjz-bot account, this is a comment-only review rather than a self-request-changes event.
The quota is about to reset, so I am concentrating the remaining token budget on these reviews.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
for more information, see https://pre-commit.ci
3 participants
Fixes #5627
Summary
nlocalandnghostwith one element instead of sizing them bynswapgpuMemcpyonly for actual GPU places and uses hostmemcpyfor CPU or host-pinned placesnswap == 0, CPU self-copy, and the reverse self-swap used by autogradWhy existing tests missed this
The existing Paddle suite did not directly call
border_op. Model-level tests therefore did not construct the two boundary conditions that matter here: a valid no-swap invocation where the atom-count scalars still need storage, and a self-swap using CPU data from an operator compiled with CUDA support. Normal multi-rank runs also tend to use tensors on the configured accelerator, hiding the mismatch between CUDA-awareness and actual tensor place.The new no-swap test passes empty communication arrays with scalar atom counts, while the self-swap test keeps a real LAMMPS-style pointer-valued send list alive and checks both forward data and backward execution on CPU tensors. In a CUDA-enabled CI build, the historical code would route that CPU pointer through device-to-device
gpuMemcpy.Validation
source/op/pd/setup.pypytest source/tests/pd/test_border_op.py -q(2 passed, including backward)GOOGLE_CUDA + USE_MPIbranch with Paddle, CUDA 12.4, and MPI headers usingmpicxx -fsyntax-onlyruff format .ruff check .clang-format --dry-run --Werror source/op/pd/comm.ccgit diff --checkThe local Paddle wheel is CPU-only, so runtime execution of the CUDA-enabled custom op is left to CUDA CI; the CUDA/MPI branch was still compiled locally, and the new CPU-place self-swap test is designed to run unchanged in that build.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
Tests